#include "tcpecho.h"
#include "lwip/netif.h"
#include "lwip/ip.h"
#include "lwip/tcp.h"
#include "lwip/init.h"
#include "netif/etharp.h"
#include "lwip/udp.h"
#include "lwip/pbuf.h"
#include <stdio.h>
#include <string.h>


//static err_t tcpecho_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err)
//{
//  if (p != NULL)
//  {
//
//  tcp_recved(tpcb, p->tot_len);
//
//
//  tcp_write(tpcb, p->payload, p->tot_len, 1);
//
//  memset(p->payload, 0 , p->tot_len);
//  pbuf_free(p);
//
//  }
//  else if (err == ERR_OK)
//  {
//    return tcp_close(tpcb);
//  }
//  return ERR_OK;
//}

//static err_t tcpecho_accept(void *arg, struct tcp_pcb *newpcb, err_t err)
//{

//  tcp_recv(newpcb, tcpecho_recv);
//  return ERR_OK;
//}

//void TCP_Echo_Init(void)
//{
//  struct tcp_pcb *pcb = NULL;
//
//
//  pcb = tcp_new();
//
//
//  tcp_bind(pcb, IP_ADDR_ANY, TCP_ECHO_PORT);


//
//  pcb = tcp_listen(pcb);

//
//  tcp_accept(pcb, tcpecho_accept);
//}


static struct tcp_pcb* tcp_echoserver_pcb;

/* ECHO protocol states */
enum tcp_echoserver_states {
    ES_NONE = 0,
    ES_ACCEPTED,
    ES_RECEIVED,
    ES_CLOSING
};

/* structure for maintaing connection infos to be passed as argument
   to LwIP callbacks*/
struct tcp_echoserver_struct {
    u8_t state;             /* current connection state */
    u8_t retries;
    struct tcp_pcb* pcb;    /* pointer on the current tcp_pcb */
    struct pbuf* p;         /* pointer on the received/to be transmitted pbuf */
};


static err_t tcp_echoserver_accept(void* arg, struct tcp_pcb* newpcb, err_t err);
static err_t tcp_echoserver_recv(void* arg, struct tcp_pcb* tpcb, struct pbuf* p, err_t err);
static void tcp_echoserver_error(void* arg, err_t err);
static err_t tcp_echoserver_poll(void* arg, struct tcp_pcb* tpcb);
static err_t tcp_echoserver_sent(void* arg, struct tcp_pcb* tpcb, u16_t len);
static void tcp_echoserver_send(struct tcp_pcb* tpcb, struct tcp_echoserver_struct* es);
static void tcp_echoserver_connection_close(struct tcp_pcb* tpcb, struct tcp_echoserver_struct* es);


/**
  * @brief  Initializes the tcp echo server
  * @param  None
  * @retval None
  */
void tcp_echoserver_init(void)
{
    /* create new tcp pcb */
    tcp_echoserver_pcb = tcp_new();

    if (tcp_echoserver_pcb != NULL) {
        err_t err;

        /* bind echo_pcb to port 7 (ECHO protocol) */
        err = tcp_bind(tcp_echoserver_pcb, IP_ADDR_ANY, TCP_ECHO_PORT);

        if (err == ERR_OK) {
            /* start tcp listening for echo_pcb */
            tcp_echoserver_pcb = tcp_listen(tcp_echoserver_pcb);
            printf("Start listening \r\n");
            /* initialize LwIP tcp_accept callback function */
            tcp_accept(tcp_echoserver_pcb, tcp_echoserver_accept);
            printf("Mount the client connection callback function \r\n");
        }
        else {
            /* deallocate the pcb */
            memp_free(MEMP_TCP_PCB, tcp_echoserver_pcb);
            printf("TCP PCB Memory request failed \r\n");
        }
    }
}

/**
  * @brief  This function is the implementation of tcp_accept LwIP callback
  * @param  arg: not used
  * @param  newpcb: pointer on tcp_pcb struct for the newly created tcp connection
  * @param  err: not used
  * @retval err_t: error status
  */
static err_t tcp_echoserver_accept(void* arg, struct tcp_pcb* newpcb, err_t err)
{
    err_t ret_err;
    struct tcp_echoserver_struct* es;

    LWIP_UNUSED_ARG(arg);
    LWIP_UNUSED_ARG(err);

    /* set priority for the newly accepted tcp connection newpcb */
    tcp_setprio(newpcb, TCP_PRIO_MIN);
    printf("Receive client connection request, set the client that just connected to the lowest priority \r\n");

    uint8_t iptxt[20];
    sprintf((char*)iptxt, "%s", ip4addr_ntoa((const ip4_addr_t*)&newpcb->remote_ip));
    printf ("Client IP address: %s\r\n", iptxt);

    /* allocate structure es to maintain tcp connection informations */
    es = (struct tcp_echoserver_struct*)mem_malloc(sizeof(struct tcp_echoserver_struct));
    if (es != NULL) {
        es->state = ES_ACCEPTED;
        es->pcb = newpcb;
        es->retries = 0;
        es->p = NULL;

        printf("Mount the required callback function for the newly connected Client \r\n");

        /* pass newly allocated es structure as argument to newpcb */
        tcp_arg(newpcb, es);

        /* initialize lwip tcp_recv callback function for newpcb  */
        tcp_recv(newpcb, tcp_echoserver_recv);

        /* initialize lwip tcp_err callback function for newpcb  */
        tcp_err(newpcb, tcp_echoserver_error);

        /* initialize lwip tcp_poll callback function for newpcb */
        tcp_poll(newpcb, tcp_echoserver_poll, 0);
        ret_err = ERR_OK;
    }
    else {
        /*  close tcp connection */
        tcp_echoserver_connection_close(newpcb, es);
        printf("tcp_echoserver_struct Memory request failed Close the connection \r\n");
        /* return memory error */
        ret_err = ERR_MEM;
    }
    return ret_err;
}


/**
  * @brief  This function is the implementation for tcp_recv LwIP callback
  * @param  arg: pointer on a argument for the tcp_pcb connection
  * @param  tpcb: pointer on the tcp_pcb connection
  * @param  pbuf: pointer on the received pbuf
  * @param  err: error information regarding the reveived pbuf
  * @retval err_t: error code
  */
static err_t tcp_echoserver_recv(void* arg, struct tcp_pcb* tpcb, struct pbuf* p, err_t err)
{
    struct tcp_echoserver_struct* es;
    err_t ret_err;

    LWIP_ASSERT("arg != NULL", arg != NULL);
    printf("Receive client data\r\n");
    es = (struct tcp_echoserver_struct*)arg;

    /* if we receive an empty tcp frame from client => close connection */
    if (p == NULL) {
        printf("Disconnect request received \r\n");
        /* remote host closed connection */
        es->state = ES_CLOSING;
        if(es->p == NULL) {
            /* we're done sending, close connection */
            tcp_echoserver_connection_close(tpcb, es);
        }
        else {
            printf("The data sent has not been sent yet \r\n");
            /* we're not done yet */
            /* acknowledge received packet */
            tcp_sent(tpcb, tcp_echoserver_sent);
            printf("Load and send completion callback function \r\n");
            /* send remaining data*/
            tcp_echoserver_send(tpcb, es);
        }
        ret_err = ERR_OK;
    }
    /* else : a non empty frame was received from client but for some reason err != ERR_OK */
    else if(err != ERR_OK) {
        /* free received pbuf*/
        if (p != NULL) {
            es->p = NULL;
            pbuf_free(p);
        }
        ret_err = err;
    }
    else if(es->state == ES_ACCEPTED) {
        /* first data chunk in p->payload */
        es->state = ES_RECEIVED;

        /* store reference to incoming pbuf (chain) */
        es->p = p;

        /* initialize LwIP tcp_sent callback function */
        tcp_sent(tpcb, tcp_echoserver_sent);
        printf("The device is just connected, mount the data just received, and set the sending completion callback\r\n");
        /* send back the received data (echo) */
        tcp_echoserver_send(tpcb, es);

        ret_err = ERR_OK;
    }
    else if (es->state == ES_RECEIVED) {
        /* more data received from client and previous data has been already sent*/
        if(es->p == NULL) {
            es->p = p;
            printf("The received data is sent back directly \r\n");
            /* send back received data */
            tcp_echoserver_send(tpcb, es);
        }
        else {
            struct pbuf* ptr;
            printf("The last data has not been sent yet, splice the new data at the back \r\n");
            /* chain pbufs to the end of what we recv'ed previously  */
            ptr = es->p;
            pbuf_chain(ptr, p);
        }
        ret_err = ERR_OK;
    }
    else if(es->state == ES_CLOSING) {
        printf("The connection is closed, but the data is still received \r\n");
        /* odd case, remote side closing twice, trash data */
        tcp_recved(tpcb, p->tot_len);
        es->p = NULL;
        pbuf_free(p);
        ret_err = ERR_OK;
    }
    else {
        /* unkown es->state, trash data  */
        tcp_recved(tpcb, p->tot_len);
        es->p = NULL;
        pbuf_free(p);
        ret_err = ERR_OK;
    }
    return ret_err;
}

/**
  * @brief  This function implements the tcp_err callback function (called
  *         when a fatal tcp_connection error occurs.
  * @param  arg: pointer on argument parameter
  * @param  err: not used
  * @retval None
  */
static void tcp_echoserver_error(void* arg, err_t err)
{
    struct tcp_echoserver_struct* es;

    LWIP_UNUSED_ARG(err);
    printf("error: %d \r\n", err);
    es = (struct tcp_echoserver_struct*)arg;
    if (es != NULL) {
        /*  free es structure */
        mem_free(es);
    }
}

/**
  * @brief  This function implements the tcp_poll LwIP callback function
  * @param  arg: pointer on argument passed to callback
  * @param  tpcb: pointer on the tcp_pcb for the current tcp connection
  * @retval err_t: error code
  */
static err_t tcp_echoserver_poll(void* arg, struct tcp_pcb* tpcb)
{
    err_t ret_err;
    struct tcp_echoserver_struct* es;

    es = (struct tcp_echoserver_struct*)arg;
    if (es != NULL) {
        if (es->p != NULL) {
            tcp_sent(tpcb, tcp_echoserver_sent);
            /* there is a remaining pbuf (chain) , try to send data */
            tcp_echoserver_send(tpcb, es);
        }
        else {
            /* no remaining pbuf (chain)  */
            if(es->state == ES_CLOSING) {
                /*  close tcp connection */
                tcp_echoserver_connection_close(tpcb, es);
            }
        }
        ret_err = ERR_OK;
    }
    else {
        /* nothing to be done */
        tcp_abort(tpcb);
        ret_err = ERR_ABRT;
    }
    return ret_err;
}

/**
  * @brief  This function implements the tcp_sent LwIP callback (called when ACK
  *         is received from remote host for sent data)
  * @param  None
  * @retval None
  */
static err_t tcp_echoserver_sent(void* arg, struct tcp_pcb* tpcb, u16_t len)
{
    struct tcp_echoserver_struct* es;

    LWIP_UNUSED_ARG(len);

    es = (struct tcp_echoserver_struct*)arg;
    es->retries = 0;

    if(es->p != NULL) {
        /* still got pbufs to send */
        tcp_sent(tpcb, tcp_echoserver_sent);
        tcp_echoserver_send(tpcb, es);
    }
    else {
        /* if no more data to send and client closed connection*/
        if(es->state == ES_CLOSING)
            tcp_echoserver_connection_close(tpcb, es);
    }
    return ERR_OK;
}


/**
  * @brief  This function is used to send data for tcp connection
  * @param  tpcb: pointer on the tcp_pcb connection
  * @param  es: pointer on echo_state structure
  * @retval None
  */
static void tcp_echoserver_send(struct tcp_pcb* tpcb, struct tcp_echoserver_struct* es)
{
    struct pbuf* ptr;
    err_t wr_err = ERR_OK;
    uint16_t Count ;
    printf("Total length of data sent : %ld \r\n", es->p->tot_len);
    printf("send data: \r\n");
    while ((wr_err == ERR_OK) &&
            (es->p != NULL) &&
            (es->p->len <= tcp_sndbuf(tpcb))) {

        /* get pointer on pbuf from es structure */
        ptr = es->p;

        for( Count = 0 ; Count < ptr->len; Count++ ) {
            printf("0X%02X ", ((uint8_t*)ptr->payload)[Count]);
        }
        /* enqueue data for transmission */
        wr_err = tcp_write(tpcb, ptr->payload, ptr->len, TCP_WRITE_FLAG_COPY);

        if (wr_err == ERR_OK) {
            u16_t plen;
            u8_t freed;

            plen = ptr->len;

            /* continue with next pbuf in chain (if any) */
            es->p = ptr->next;

            if(es->p != NULL) {
                /* increment reference count for es->p */
                pbuf_ref(es->p);
            }

            /* chop first pbuf from chain */
            do {
                /* try hard to free pbuf */
                freed = pbuf_free(ptr);
            } while(freed == 0);
            /* we can read more data now */
            tcp_recved(tpcb, plen);
        }
        else if(wr_err == ERR_MEM) {
            /* we are low on memory, try later / harder, defer to poll */
            es->p = ptr;
        }
        else {

        }
    }
    printf("\r\n");
}

/**
  * @brief  This functions closes the tcp connection
  * @param  tcp_pcb: pointer on the tcp connection
  * @param  es: pointer on echo_state structure
  * @retval None
  */
static void tcp_echoserver_connection_close(struct tcp_pcb* tpcb, struct tcp_echoserver_struct* es)
{

    /* remove all callbacks */
    tcp_arg(tpcb, NULL);
    tcp_sent(tpcb, NULL);
    tcp_recv(tpcb, NULL);
    tcp_err(tpcb, NULL);
    tcp_poll(tpcb, NULL, 0);

    /* delete es structure */
    if (es != NULL) {
        mem_free(es);
    }

    /* close tcp connection */
    tcp_close(tpcb);
    printf("Connection closed \r\n");
}






